[Tune] Enable experiment restore from moved cloud uri #31669

justinvyu · 2023-01-13T19:34:02Z

Why are these changes needed?

This PR fixes experiment restoration from a different cloud URI to save future results and checkpoints to new URI instead of continuing to write to the old location. The workflow of starting a local experiment, uploading the experiment dir to cloud, then restoring from the URI from a different cluster is also possible now.

Example:

# initial experiment experiment directory at "s3://original/exp_name"
# experiment directory moved to "s3://new/new_exp_name"
Tuner.restore("s3://new/new_exp_name")
# new results/checkpoints should get synced to the new URI "s3://new/new_exp_name"

This is a follow-up to #29920, which enabled experiment restoration from a moved local experiment directory.

I also took the chance to simplify some of the trial loading/serialization logic.

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: Justin Yu <justinvyu@berkeley.edu> Fix constructor args Signed-off-by: Justin Yu <justinvyu@berkeley.edu> Fix all references to _checkpoint_dir Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

…e ckpt dir Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

justinvyu · 2023-01-13T19:37:18Z

python/ray/tune/execution/trial_runner.py

+            # since the dir might not be creatable locally.
+            # TODO(ekl) this is kind of a hack.
+            if not ray.util.client.ray.is_connected():
+                trial.init_logdir()  # Create logdir if it does not exist


I moved this out of Trial.__setstate__. Are trials serialized/deserialized in any situation that would require this to stay within __setstate__? My understanding is that trial objects don't get shipped around and stay on the Tune driver. Only serialized/deserialized on experiment checkpoint and restore, which is handled here.

What's the reason for moving this?

You're correct @justinvyu and I think your solution is much cleaner. For setstate/getstate, it's ok to just restore exactly the state that was saved. If properties are overwritten, that should happen in the function that issues the restore. No need for magic here.

@krfricke doesn't this make trials effectively unserializable in the general case unless the trainable is registered? I don't think that's an issue, but perhaps something to consider

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

…ore_from_moved_cloud_uri

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

…m trial runner tests Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

…ore_from_moved_cloud_uri

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

krfricke

This is great, thanks! Only a couple of nits

python/ray/tune/analysis/experiment_analysis.py

krfricke · 2023-01-24T01:12:26Z

python/ray/tune/execution/trial_runner.py

+            # since the dir might not be creatable locally.
+            # TODO(ekl) this is kind of a hack.
+            if not ray.util.client.ray.is_connected():
+                trial.init_logdir()  # Create logdir if it does not exist


You're correct @justinvyu and I think your solution is much cleaner. For setstate/getstate, it's ok to just restore exactly the state that was saved. If properties are overwritten, that should happen in the function that issues the restore. No need for magic here.

python/ray/tune/tests/test_tuner_restore.py

…ore_from_moved_cloud_uri

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

…2010) #31669 changed the `Trial.__dict__` by moving `local_dir` to `_local_dir`, which resulted in an error in our tune cloud tests. This PR updates the signature of the `TrialStub` class to resolve the issue. Signed-off-by: Kai Fricke <kai@anyscale.com>

…y-project#32010) ray-project#31669 changed the `Trial.__dict__` by moving `local_dir` to `_local_dir`, which resulted in an error in our tune cloud tests. This PR updates the signature of the `TrialStub` class to resolve the issue. Signed-off-by: Kai Fricke <kai@anyscale.com> Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>

…th (#40647) This is a known regression introduced in 2.7: moving the path of the experiment directory and attempting to restore the experiment and/or the experiment results doesn't work due to the absolute paths saved in the trial metadata. This PR implements a fix similar to #31669 -- replacing the root of the tracked checkpoint paths with the new storage path, and updating on experiment restoration / result loading from a path. --------- Signed-off-by: Justin Yu <justinvyu@anyscale.com>

justinvyu added 9 commits January 13, 2023 11:16

Remove duplicate local_dir + rename to be the same as trial runner

97181c9

Signed-off-by: Justin Yu <justinvyu@berkeley.edu> Fix constructor args Signed-off-by: Justin Yu <justinvyu@berkeley.edu> Fix all references to _checkpoint_dir Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Add Trial.from_json_state alternate constructor

9f38a34

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Only save around sync config + experiment dir name to construct remot…

af2b189

…e ckpt dir Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Update upload dir and experiment name upon URI restore

d225599

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Refactor relative checkpoint dir logic to be in local_dir setter instead

e0267a8

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Refactor trial loading logic to use Trial.from_json_state

87c5452

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Switch ExperimentAnalysis trial loading to from_json_state as well

e202dc0

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Add restore test from moved URI

dde5d87

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Remove load trial utilities

f6299a3

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

justinvyu added tune Tune-related issues air labels Jan 13, 2023

justinvyu commented Jan 13, 2023

View reviewed changes

justinvyu added 14 commits January 16, 2023 10:17

Fix sync_config=None case

7f810bd

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Fix upload dir parsing to include query string

936174c

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Merge branch 'master' of https://github.com/ray-project/ray into rest…

92729aa

…ore_from_moved_cloud_uri

Sort trials in results by trial id

20104dc

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Strict input type for from_json_state

bce61b4

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Clear memory filesys fixture for tests that use it

3ca887d

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Remove useless check (trial logdir is loaded from the checkpoint)

12068a8

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Fix tests related to from_json_state

687ea5e

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Update TrialRunner docstring, remove remote_checkpoint_dir kwargs fro…

c65999f

…m trial runner tests Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Simplify mocks in test_api

f16e513

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Merge branch 'master' of https://github.com/ray-project/ray into rest…

3e5bb5b

…ore_from_moved_cloud_uri

Remove unnecessary import

69a7221

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Remove really old backwards compatibility test

fdeaa19

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Remove unnecessary import

b54c30c

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

justinvyu assigned krfricke Jan 18, 2023

justinvyu requested review from krfricke, Yard1 and xwjiang2010 January 18, 2023 17:58

justinvyu marked this pull request as ready for review January 18, 2023 17:58

krfricke approved these changes Jan 24, 2023

View reviewed changes

justinvyu added 3 commits January 23, 2023 22:21

Merge branch 'master' of https://github.com/ray-project/ray into rest…

b814b72

…ore_from_moved_cloud_uri

Remove unused decode fn

2132cfb

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

Clarify why the checkpoint gets moved

3257b9a

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>

krfricke merged commit 47f4da3 into ray-project:master Jan 24, 2023

krfricke mentioned this pull request Jan 27, 2023

[tune] Fix tune_cloud_* tests fow new Trial dict #32010

Merged

7 tasks

justinvyu deleted the restore_from_moved_cloud_uri branch April 26, 2023 20:13

justinvyu mentioned this pull request Aug 8, 2023

[air] pyarrow.fs persistence (9/n): ray.train.Checkpoint restore: Manual restore #38128

Merged

11 tasks

justinvyu mentioned this pull request Oct 24, 2023

[tune/train] Restore Tuner and results properly from moved storage path #40647

Merged

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Tune] Enable experiment restore from moved cloud uri #31669

[Tune] Enable experiment restore from moved cloud uri #31669

justinvyu commented Jan 13, 2023

justinvyu Jan 13, 2023

Yard1 Jan 23, 2023

krfricke Jan 24, 2023

Yard1 Jan 24, 2023

krfricke left a comment

krfricke Jan 24, 2023

[Tune] Enable experiment restore from moved cloud uri #31669

[Tune] Enable experiment restore from moved cloud uri #31669

Conversation

justinvyu commented Jan 13, 2023

Why are these changes needed?

Related issue number

Checks

justinvyu Jan 13, 2023

Choose a reason for hiding this comment

Yard1 Jan 23, 2023

Choose a reason for hiding this comment

krfricke Jan 24, 2023

Choose a reason for hiding this comment

Yard1 Jan 24, 2023

Choose a reason for hiding this comment

krfricke left a comment

Choose a reason for hiding this comment

krfricke Jan 24, 2023

Choose a reason for hiding this comment